import torch
import jieba
from transformers import AutoTokenizer, AutoModel

# 加载预训练的分词器和模型
tokenizer = AutoTokenizer.from_pretrained("bert-base-chinese")
model = AutoModel.from_pretrained("bert-base-chinese")

# 示例中文文本
text = "这是一个测试的句子。"

# 使用 jieba 进行分词
words = jieba.lcut(text)
print("分词结果:", words)

# 存储每个词的向量
word_vectors = []

# 对每个词进行编码并获取向量
for word in words:
    # 对单个词进行编码
    inputs = tokenizer(word, return_tensors='pt')
    with torch.no_grad():
        outputs = model(**inputs)
    # 取 [CLS] 标记之后的第一个词向量作为该词的表示
    word_vector = outputs.last_hidden_state[:, 1:-1, :].mean(dim=1).squeeze()
    word_vectors.append(word_vector)

# 将每个词的向量组合成一个张量
sentence_tensor = torch.stack(word_vectors, dim=0)

print("句子张量的形状:", sentence_tensor.shape)
print("句子张量的形状:", sentence_tensor)